#!/bin/bash

#==============================================================================
# 项目结构生成器 (Project Tree Generator)
#==============================================================================
# 
# 功能描述:
#   生成项目目录结构的可视化树形图，支持多种输出格式
#   可用于文档编写、项目展示、代码审查等场景
#
# 作者: Pedro.Pei
# 版本: 1.0
# 更新日期: 2025
#
# 主要功能:
#   🌳 生成ASCII树形结构 - 经典的文本树形图
#   📋 生成Markdown格式 - 适合文档和README
#   📊 生成JSON格式 - 便于程序处理
#   🎨 生成HTML格式 - 可在浏览器中查看
#   🔍 智能过滤功能 - 排除不需要的文件和目录
#   📏 深度控制 - 限制显示层级
#   📊 统计信息 - 显示文件和目录数量
#
# 使用方法:
#   ./project_tree.sh [目录路径] [选项]
#
# 选项:
#   -f, --format <格式>     输出格式 (ascii|markdown|json|html) 默认: ascii
#   -o, --output <文件>     输出到文件而不是控制台
#   -d, --depth <层级>      最大显示深度 默认: 无限制
#   -s, --size             显示文件大小
#   --include-hidden       包含隐藏文件和目录
#   --exclude <模式>       排除文件模式 (如: *.log,node_modules)
#   --only-dirs           只显示目录
#   --stats               显示统计信息
#   -c, --color           使用彩色输出 (仅ASCII格式)
#   -v, --verbose         显示详细处理过程
#   -h, --help           显示帮助信息
#
#==============================================================================

# 脚本配置
SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd)
SCRIPT_NAME=$(basename "$0")

# 默认参数
TARGET_DIR="."
OUTPUT_FORMAT="ascii"
OUTPUT_FILE=""
MAX_DEPTH=""
SHOW_SIZE=false
INCLUDE_HIDDEN=false
EXCLUDE_PATTERNS=""
ONLY_DIRS=false
SHOW_STATS=false
USE_COLOR=false
VERBOSE=false

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
CYAN='\033[0;36m'
WHITE='\033[1;37m'
GRAY='\033[0;90m'
NC='\033[0m' # No Color

# 统计变量
total_files=0
total_dirs=0
total_size=0

# 显示帮助信息
show_help() {
    echo -e "${CYAN}项目结构生成器 v1.0${NC}"
    echo ""
    echo "用法: $0 [目录路径] [选项]"
    echo ""
    echo -e "${WHITE}选项:${NC}"
    echo -e "  ${GREEN}-f, --format <格式>${NC}     输出格式 (ascii|markdown|json|html) 默认: ascii"
    echo -e "  ${GREEN}-o, --output <文件>${NC}     输出到文件而不是控制台"
    echo -e "  ${GREEN}-d, --depth <层级>${NC}      最大显示深度 默认: 无限制"
    echo -e "  ${GREEN}-s, --size${NC}             显示文件大小"
    echo -e "  ${GREEN}--include-hidden${NC}       包含隐藏文件和目录"
    echo -e "  ${GREEN}--exclude <模式>${NC}       排除文件模式 (如: *.log,node_modules)"
    echo -e "  ${GREEN}--only-dirs${NC}           只显示目录"
    echo -e "  ${GREEN}--stats${NC}               显示统计信息"
    echo -e "  ${GREEN}-c, --color${NC}           使用彩色输出 (仅ASCII格式)"
    echo -e "  ${GREEN}-v, --verbose${NC}         显示详细处理过程"
    echo -e "  ${GREEN}-h, --help${NC}           显示此帮助信息"
    echo ""
    echo -e "${WHITE}输出格式:${NC}"
    echo -e "  ${GREEN}ascii${NC}     - 经典的ASCII树形结构"
    echo -e "  ${GREEN}markdown${NC}  - Markdown格式，适合文档"
    echo -e "  ${GREEN}json${NC}      - JSON格式，便于程序处理"
    echo -e "  ${GREEN}html${NC}      - HTML格式，可在浏览器查看"
    echo ""
    echo -e "${WHITE}示例:${NC}"
    echo "  $0                                    # 显示当前目录结构"
    echo "  $0 ~/Projects/myapp                  # 显示指定目录结构"
    echo "  $0 -f markdown -o README.md          # 生成Markdown格式到文件"
    echo "  $0 -d 3 --exclude node_modules       # 限制深度并排除node_modules"
    echo "  $0 --only-dirs -c --stats            # 只显示目录，使用彩色和统计"
}

# 解析命令行参数
parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_help
                exit 0
                ;;
            -f|--format)
                OUTPUT_FORMAT="$2"
                if [[ ! "$OUTPUT_FORMAT" =~ ^(ascii|markdown|json|html)$ ]]; then
                    echo -e "${RED}错误: 不支持的输出格式 '$OUTPUT_FORMAT'${NC}"
                    exit 1
                fi
                shift 2
                ;;
            -o|--output)
                OUTPUT_FILE="$2"
                shift 2
                ;;
            -d|--depth)
                MAX_DEPTH="$2"
                if ! [[ "$MAX_DEPTH" =~ ^[0-9]+$ ]]; then
                    echo -e "${RED}错误: depth参数必须是数字${NC}"
                    exit 1
                fi
                shift 2
                ;;
            -s|--size)
                SHOW_SIZE=true
                shift
                ;;
            --include-hidden)
                INCLUDE_HIDDEN=true
                shift
                ;;
            --exclude)
                EXCLUDE_PATTERNS="$2"
                shift 2
                ;;
            --only-dirs)
                ONLY_DIRS=true
                shift
                ;;
            --stats)
                SHOW_STATS=true
                shift
                ;;
            -c|--color)
                USE_COLOR=true
                shift
                ;;
            -v|--verbose)
                VERBOSE=true
                shift
                ;;
            -*)
                echo -e "${RED}错误: 未知选项 '$1'${NC}"
                show_help
                exit 1
                ;;
            *)
                if [ -z "$TARGET_DIR" ] || [ "$TARGET_DIR" = "." ]; then
                    TARGET_DIR="$1"
                else
                    echo -e "${RED}错误: 只能指定一个目录${NC}"
                    exit 1
                fi
                shift
                ;;
        esac
    done

    # 验证目录存在
    if [ ! -d "$TARGET_DIR" ]; then
        echo -e "${RED}错误: 目录不存在 '$TARGET_DIR'${NC}"
        exit 1
    fi

    # 转换为绝对路径
    if ! TARGET_DIR=$(cd "$TARGET_DIR" 2>/dev/null && pwd); then
        echo -e "${RED}错误: 无法访问目录 '$TARGET_DIR'${NC}"
        exit 1
    fi
    
    # 如果不是ASCII格式，禁用颜色
    if [ "$OUTPUT_FORMAT" != "ascii" ]; then
        USE_COLOR=false
    fi
}

# 格式化文件大小
format_size() {
    local bytes=$1
    if [ "$bytes" -eq 0 ]; then
        echo "0B"
    elif [ "$bytes" -lt 1024 ]; then
        echo "${bytes}B"
    elif [ "$bytes" -lt 1048576 ]; then
        echo "$(echo "scale=1; $bytes/1024" | bc)K"
    elif [ "$bytes" -lt 1073741824 ]; then
        echo "$(echo "scale=1; $bytes/1048576" | bc)M"
    else
        echo "$(echo "scale=1; $bytes/1073741824" | bc)G"
    fi
}

# 检查是否应该排除文件/目录
should_exclude() {
    local name="$1"
    
    # 检查隐藏文件
    if [ "$INCLUDE_HIDDEN" = false ] && [[ "$name" == .* ]]; then
        return 0  # 排除
    fi
    
    # 检查排除模式
    if [ -n "$EXCLUDE_PATTERNS" ]; then
        IFS=',' read -ra PATTERNS <<< "$EXCLUDE_PATTERNS"
        for pattern in "${PATTERNS[@]}"; do
            if [[ "$name" == $pattern ]]; then
                return 0  # 排除
            fi
        done
    fi
    
    return 1  # 不排除
}

# 获取彩色文件名
get_colored_name() {
    local path="$1"
    local name=$(basename "$path")
    
    # 如果输出到文件或者不使用颜色，则返回普通名称
    if [ "$USE_COLOR" = false ] || [ -n "$OUTPUT_FILE" ]; then
        echo "$name"
        return
    fi
    
    if [ -d "$path" ]; then
        echo -e "${BLUE}$name${NC}"
    elif [ -x "$path" ]; then
        echo -e "${GREEN}$name${NC}"
    elif [[ "$name" == *.md ]]; then
        echo -e "${CYAN}$name${NC}"
    elif [[ "$name" == *.json ]] || [[ "$name" == *.xml ]] || [[ "$name" == *.yaml ]] || [[ "$name" == *.yml ]]; then
        echo -e "${YELLOW}$name${NC}"
    elif [[ "$name" == *.sh ]] || [[ "$name" == *.py ]] || [[ "$name" == *.js ]] || [[ "$name" == *.ts ]]; then
        echo -e "${GREEN}$name${NC}"
    elif [[ "$name" == *.txt ]] || [[ "$name" == *.log ]]; then
        echo -e "${GRAY}$name${NC}"
    else
        echo "$name"
    fi
}

# 生成ASCII格式树
generate_ascii_tree() {
    local dir="$1"
    local prefix="$2"
    local current_depth="$3"
    
    # 检查深度限制
    if [ -n "$MAX_DEPTH" ] && [ "$current_depth" -gt "$MAX_DEPTH" ]; then
        return
    fi
    
    if [ "$VERBOSE" = true ]; then
        echo -e "${GRAY}处理目录: $dir${NC}" >&2
    elif [ "$current_depth" -le 2 ]; then
        # 对于浅层目录，即使非verbose模式也显示进度
        local relative_path=$(echo "$dir" | sed "s|$TARGET_DIR|.|")
        echo -e "${GRAY}正在处理: $relative_path${NC}" >&2
    fi
    
    # 获取目录内容并排序
    local items=()
    while IFS= read -r -d '' item; do
        local name=$(basename "$item")
        if ! should_exclude "$name"; then
            if [ "$ONLY_DIRS" = true ]; then
                [ -d "$item" ] && items+=("$item")
            else
                items+=("$item")
            fi
        fi
    done < <(find "$dir" -maxdepth 1 -mindepth 1 -print0 | sort -z)
    
    local count=${#items[@]}
    local i=0
    
    for item in "${items[@]}"; do
        i=$((i + 1))
        local name=$(basename "$item")
        local is_last=$((i == count))
        
        # 选择前缀符号
        if [ "$is_last" = 1 ]; then
            if [ -n "$OUTPUT_FILE" ]; then
                # 输出到文件时使用ASCII字符
                local current_prefix="\`-- "
                local next_prefix="$prefix    "
            else
                # 终端输出时使用UTF-8字符
                local current_prefix="└── "
                local next_prefix="$prefix    "
            fi
        else
            if [ -n "$OUTPUT_FILE" ]; then
                # 输出到文件时使用ASCII字符
                local current_prefix="|-- "
                local next_prefix="$prefix|   "
            else
                # 终端输出时使用UTF-8字符
                local current_prefix="├── "
                local next_prefix="$prefix│   "
            fi
        fi
        
        # 获取彩色名称
        local colored_name=$(get_colored_name "$item")
        
        # 添加大小信息
        local size_info=""
        if [ "$SHOW_SIZE" = true ] && [ -f "$item" ]; then
            local file_size=$(stat -f%z "$item" 2>/dev/null || echo 0)
            size_info=" ($(format_size $file_size))"
            total_size=$((total_size + file_size))
        fi
        
        # 输出当前项
        echo "${prefix}${current_prefix}${colored_name}${size_info}"
        
        # 更新统计
        if [ -d "$item" ]; then
            total_dirs=$((total_dirs + 1))
            # 递归处理子目录
            generate_ascii_tree "$item" "$next_prefix" $((current_depth + 1))
        else
            total_files=$((total_files + 1))
        fi
    done
}

# 生成Markdown格式树
generate_markdown_tree() {
    local dir="$1"
    local prefix="$2"
    local current_depth="$3"
    
    # 检查深度限制
    if [ -n "$MAX_DEPTH" ] && [ "$current_depth" -gt "$MAX_DEPTH" ]; then
        return
    fi
    
    # 获取目录内容并排序
    local items=()
    while IFS= read -r -d '' item; do
        local name=$(basename "$item")
        if ! should_exclude "$name"; then
            if [ "$ONLY_DIRS" = true ]; then
                [ -d "$item" ] && items+=("$item")
            else
                items+=("$item")
            fi
        fi
    done < <(find "$dir" -maxdepth 1 -mindepth 1 -print0 | sort -z)
    
    for item in "${items[@]}"; do
        local name=$(basename "$item")
        
        # 添加大小信息
        local size_info=""
        if [ "$SHOW_SIZE" = true ] && [ -f "$item" ]; then
            local file_size=$(stat -f%z "$item" 2>/dev/null || echo 0)
            size_info=" *($(format_size $file_size))*"
            total_size=$((total_size + file_size))
        fi
        
        # 输出当前项
        if [ -d "$item" ]; then
            echo "${prefix}- **${name}/**${size_info}"
            total_dirs=$((total_dirs + 1))
            # 递归处理子目录
            generate_markdown_tree "$item" "$prefix  " $((current_depth + 1))
        else
            echo "${prefix}- ${name}${size_info}"
            total_files=$((total_files + 1))
        fi
    done
}

# 生成JSON格式树
generate_json_tree() {
    local dir="$1"
    local current_depth="$2"
    local is_root="${3:-true}"
    
    # 检查深度限制
    if [ -n "$MAX_DEPTH" ] && [ "$current_depth" -gt "$MAX_DEPTH" ]; then
        echo "null"
        return
    fi
    
    local json_output=""
    if [ "$is_root" = "true" ]; then
        json_output="{\n  \"name\": \"$(basename "$dir")\",\n  \"type\": \"directory\",\n  \"path\": \"$dir\",\n  \"children\": ["
    else
        json_output="["
    fi
    
    # 获取目录内容并排序
    local items=()
    while IFS= read -r -d '' item; do
        local name=$(basename "$item")
        if ! should_exclude "$name"; then
            if [ "$ONLY_DIRS" = true ]; then
                [ -d "$item" ] && items+=("$item")
            else
                items+=("$item")
            fi
        fi
    done < <(find "$dir" -maxdepth 1 -mindepth 1 -print0 | sort -z)
    
    local count=${#items[@]}
    local i=0
    
    for item in "${items[@]}"; do
        i=$((i + 1))
        local name=$(basename "$item")
        local is_last=$((i == count))
        
        if [ -d "$item" ]; then
            total_dirs=$((total_dirs + 1))
            json_output+="\n    {\n      \"name\": \"$name\",\n      \"type\": \"directory\",\n      \"path\": \"$item\""
            
            # 递归获取子目录
            local children=$(generate_json_tree "$item" $((current_depth + 1)) "false")
            if [ "$children" != "[]" ]; then
                json_output+=",\n      \"children\": $children"
            fi
            json_output+="\n    }"
        else
            total_files=$((total_files + 1))
            local file_size=0
            if [ "$SHOW_SIZE" = true ]; then
                file_size=$(stat -f%z "$item" 2>/dev/null || echo 0)
                total_size=$((total_size + file_size))
            fi
            
            json_output+="\n    {\n      \"name\": \"$name\",\n      \"type\": \"file\",\n      \"path\": \"$item\""
            if [ "$SHOW_SIZE" = true ]; then
                json_output+=",\n      \"size\": $file_size"
            fi
            json_output+="\n    }"
        fi
        
        if [ "$is_last" != 1 ]; then
            json_output+=","
        fi
    done
    
    json_output+="\n  ]"
    
    if [ "$is_root" = "true" ]; then
        json_output+="\n}"
    fi
    
    echo -e "$json_output"
}

# 生成HTML格式树
generate_html_tree() {
    local dir="$1"
    local current_depth="$2"
    local is_root="${3:-true}"
    
    if [ "$is_root" = "true" ]; then
        cat << EOF
<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>项目结构 - $(basename "$dir")</title>
    <style>
        body { font-family: 'Courier New', monospace; margin: 20px; background: #f5f5f5; }
        .tree { background: white; padding: 20px; border-radius: 8px; box-shadow: 0 2px 10px rgba(0,0,0,0.1); }
        .tree ul { list-style: none; margin: 0; padding-left: 20px; }
        .tree li { margin: 2px 0; position: relative; }
        .tree li:before { content: '├── '; color: #666; }
        .tree li:last-child:before { content: '└── '; }
        .directory { color: #0066cc; font-weight: bold; }
        .file { color: #333; }
        .executable { color: #00aa00; }
        .size { color: #888; font-size: 0.9em; }
        .stats { margin-top: 20px; padding: 15px; background: #e8f4f8; border-radius: 5px; }
        h1 { color: #333; border-bottom: 2px solid #0066cc; padding-bottom: 10px; }
    </style>
</head>
<body>
    <div class="tree">
        <h1>📁 项目结构: $(basename "$dir")</h1>
        <ul>
EOF
    fi
    
    # 检查深度限制
    if [ -n "$MAX_DEPTH" ] && [ "$current_depth" -gt "$MAX_DEPTH" ]; then
        return
    fi
    
    # 获取目录内容并排序
    local items=()
    while IFS= read -r -d '' item; do
        local name=$(basename "$item")
        if ! should_exclude "$name"; then
            if [ "$ONLY_DIRS" = true ]; then
                [ -d "$item" ] && items+=("$item")
            else
                items+=("$item")
            fi
        fi
    done < <(find "$dir" -maxdepth 1 -mindepth 1 -print0 | sort -z)
    
    for item in "${items[@]}"; do
        local name=$(basename "$item")
        
        if [ -d "$item" ]; then
            total_dirs=$((total_dirs + 1))
            echo "            <li><span class=\"directory\">$name/</span>"
            echo "                <ul>"
            generate_html_tree "$item" $((current_depth + 1)) "false"
            echo "                </ul>"
            echo "            </li>"
        else
            total_files=$((total_files + 1))
            local class="file"
            if [ -x "$item" ]; then
                class="executable"
            fi
            
            local size_info=""
            if [ "$SHOW_SIZE" = true ]; then
                local file_size=$(stat -f%z "$item" 2>/dev/null || echo 0)
                size_info=" <span class=\"size\">($(format_size $file_size))</span>"
                total_size=$((total_size + file_size))
            fi
            
            echo "            <li><span class=\"$class\">$name</span>$size_info</li>"
        fi
    done
    
    if [ "$is_root" = "true" ]; then
        cat << EOF
        </ul>
EOF
        if [ "$SHOW_STATS" = true ]; then
            cat << EOF
        <div class="stats">
            <h3>📊 统计信息</h3>
            <p><strong>目录数量:</strong> $total_dirs</p>
            <p><strong>文件数量:</strong> $total_files</p>
EOF
            if [ "$SHOW_SIZE" = true ]; then
                echo "            <p><strong>总大小:</strong> $(format_size $total_size)</p>"
            fi
            echo "            <p><strong>生成时间:</strong> $(date)</p>"
            echo "        </div>"
        fi
        cat << EOF
    </div>
</body>
</html>
EOF
    fi
}

# 显示统计信息
show_statistics() {
    if [ "$SHOW_STATS" = true ]; then
        echo "" >&2
        echo -e "${CYAN}📊 统计信息${NC}" >&2
        echo -e "${CYAN}===================${NC}" >&2
        echo -e "${WHITE}目录数量:${NC} $total_dirs" >&2
        echo -e "${WHITE}文件数量:${NC} $total_files" >&2
        if [ "$SHOW_SIZE" = true ]; then
            echo -e "${WHITE}总大小:${NC} $(format_size $total_size)" >&2
        fi
        echo -e "${WHITE}生成时间:${NC} $(date)" >&2
    fi
}

# 主函数
main() {
    local start_time=$(date +%s)
    echo -e "${CYAN}🌳 项目结构生成器 v1.0${NC}" >&2
    
    if [ "$VERBOSE" = true ]; then
        echo -e "${YELLOW}目标目录: $TARGET_DIR${NC}" >&2
        echo -e "${YELLOW}输出格式: $OUTPUT_FORMAT${NC}" >&2
        echo -e "${YELLOW}开始时间: $(date)${NC}" >&2
        echo "" >&2
    fi
    
    # 重置统计
    total_files=0
    total_dirs=0
    total_size=0
    
    echo -e "${GRAY}正在生成项目结构...${NC}" >&2
    
    # 生成输出
    local output=""
    case "$OUTPUT_FORMAT" in
        "ascii")
            if [ -n "$OUTPUT_FILE" ]; then
                # 输出到文件时不使用颜色
                echo "$(basename "$TARGET_DIR")/"
            else
                # 终端输出时使用颜色
                echo -e "${BLUE}$(basename "$TARGET_DIR")/${NC}"
            fi
            generate_ascii_tree "$TARGET_DIR" "" 1
            ;;
        "markdown")
            echo "# 📁 项目结构: $(basename "$TARGET_DIR")"
            echo ""
            echo "\`\`\`"
            echo "$(basename "$TARGET_DIR")/"
            generate_markdown_tree "$TARGET_DIR" "" 1
            echo "\`\`\`"
            ;;
        "json")
            generate_json_tree "$TARGET_DIR" 1
            ;;
        "html")
            generate_html_tree "$TARGET_DIR" 1
            ;;
    esac
    
    # 显示统计信息 (除了JSON和HTML格式)
    if [[ "$OUTPUT_FORMAT" != "json" && "$OUTPUT_FORMAT" != "html" ]]; then
        show_statistics
    fi
    
    # 显示总耗时
    local end_time=$(date +%s)
    local duration=$((end_time - start_time))
    if [ "$VERBOSE" = true ] || [ "$duration" -gt 3 ]; then
        echo -e "${GREEN}✅ 生成完成，耗时: ${duration}秒${NC}" >&2
    fi
}

# 检查依赖
check_dependencies() {
    local missing_deps=()
    
    if ! command -v bc &> /dev/null; then
        missing_deps+=("bc")
    fi
    
    if ! command -v find &> /dev/null; then
        missing_deps+=("find")
    fi
    
    if [ ${#missing_deps[@]} -gt 0 ]; then
        echo -e "${RED}错误: 缺少必需的依赖:${NC}"
        for dep in "${missing_deps[@]}"; do
            echo -e "  ${YELLOW}- $dep${NC}"
        done
        exit 1
    fi
}

# 脚本入口
if [ "${BASH_SOURCE[0]}" = "${0}" ]; then
    # 检查依赖
    check_dependencies
    
    # 解析参数
    parse_arguments "$@"
    
    # 执行主函数并输出
    if [ -n "$OUTPUT_FILE" ]; then
        main > "$OUTPUT_FILE"
        echo -e "${GREEN}✅ 项目结构已生成到: $OUTPUT_FILE${NC}"
        if [ "$SHOW_STATS" = true ]; then
            show_statistics
        fi
    else
        main
    fi
fi